import numpy as np
import random
import cv2
import math

class moscia(object):
    def __init__(self, image_shape=[416,416]):
        self.targey_image_h = image_shape[0]
        self.targey_image_w = image_shape[1]

    def merge_bboxes(self, boxes_list, line_x, line_y):
        new_box = []
        for i in range(len(boxes_list)):
            for box in boxes_list[i]:
                x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
                if y2<=y1:
                    temp = y1
                    y1 = y2
                    y2 = temp
                if i == 0:      # 第一区
                    if y1 > line_y or x1 > line_x:
                        continue
                    if y2 >= line_y:
                        y2 = line_y
                    if x2 >= line_x:
                        x2 = line_x
                if i == 1:      # 第3区
                    y1=y1+line_y
                    y2 = y2 + line_y
                    if y2 < line_y or x1 > line_x:
                        continue
                    if y1 <= line_y:
                        y1 = line_y
                    if y2 >=self.targey_image_h:
                        y2=self.targey_image_h
                    if x2 >= line_x:
                        x2 = line_x
                if i == 2:  # 第4区
                    x1 = x1 + line_x
                    x2 = x2 + line_x
                    y1 = y1 + line_y
                    y2 = y2 + line_y
                    if y2 < line_y or x2 < line_x:
                        continue
                    if x1 <= line_x:
                        x1 = line_x
                    if x2>=self.targey_image_w:
                        x2=self.targey_image_w
                    if y1 <= line_y:
                        y1 = line_y
                    if y2>=self.targey_image_h:
                        y2=self.targey_image_h
                if i == 3:      # 第2区
                    x1 = x1 + line_x
                    x2 = x2 + line_x
                    if y1 > line_y or x2 < line_x:
                        continue
                    if y2 >= line_y:
                        y2 = line_y
                    if x1 <= line_x:
                        x1 = line_x
                    if x2 >=self.targey_image_w:
                        x2=self.targey_image_w
                if x2 - x1 < 40 or y2 - y1 < 40:
                    continue
                new_box.append([x1,y1,x2,y2,box[4]])
        return new_box

    def __call__(self, img_list, boxes_list):
        dx = int(self.targey_image_w * 0.4)
        dy = int(self.targey_image_h * 0.4)
        line_x = np.random.randint(dx, 416 - dx)
        line_y = np.random.randint(dy, 416 - dy)
        new_image = np.zeros([self.targey_image_h, self.targey_image_w, 3])
        new_image[:line_y, :line_x, :] = img_list[0][:line_y, :line_x, :]  # 第1区
        new_image[line_y:, :line_x, :] = img_list[1][:416 - line_y, :line_x, :]  # 第3区
        new_image[line_y:, line_x:, :] = img_list[2][:416 - line_y, :416 - line_x, :]  # 第4区
        new_image[:line_y, line_x:, :] = img_list[3][:line_y, :416 - line_x, :]  # 第2区

        new_box = self.merge_bboxes(boxes_list, line_x, line_y)

        return np.array(new_image,dtype=np.uint8), np.array(new_box,dtype=np.float)

class hsv_chance(object):
    '''hsv颜色变换'''
    def __init__(self, use=True):
        self.use = use
        self.p = 0.5
        self.hue = 0.1
        self.sat = 1.5
        self.val = 1.5

    def rand(self, a=0, b=1):
        return random.random() * (b - a) + a

    def rand_scale(self,s):
        scale = self.rand(1, s)
        if random.randint(0, 1) % 2:
            return scale
        return 1. / scale

    def __call__(self, img, bboxes):
        if self.use == False:
            return img,bboxes
        else:
            if random.random() < self.p:
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 通道切换，从而实现颜色变换
                hue = self.rand(-self.hue, self.hue)
                sat = self.rand_scale(self.sat)
                val = self.rand_scale(self.val)
                x = cv2.cvtColor(img.astype(np.float32) ,cv2.COLOR_RGB2HSV)
                x[..., 0] += hue * 360
                x[..., 0][x[..., 0] > 1] -= 1
                x[..., 0][x[..., 0] < 0] += 1
                x[..., 1] *= sat
                x[..., 2] *= val
                x[x[:, :, 0] > 360, 0] = 360
                x[:, :, 1:][x[:, :, 1:] > 1] = 1
                x[x < 0] = 0
                # img = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)
                # img = np.array(img * 255, dtype=np.uint8)     # 感觉过程都一样，下面的效果感觉好一些
                img = np.clip(cv2.cvtColor(img, cv2.COLOR_HSV2RGB), 0, 255)
            return img,bboxes

class RandomHorizontalFilp(object):
    '''水平翻转'''
    def __init__(self, use=True):
        self.use = use
        self.p =0.5

    def __call__(self, img, bboxes):
        if self.use==False:
            return img, bboxes
        else:
            if random.random() < self.p:
                _, w_img, _ = img.shape
                img = img[:, ::-1, :]
                bboxes[:, [0, 2]] = w_img - bboxes[:, [2, 0]]
            return img, bboxes

class RandomVerticalFilp(object):
    '''垂直翻转'''
    def __init__(self, use=True):
        self.use = use
        self.p =0.5

    def __call__(self, img, bboxes):
        if self.use==False:
            return img, bboxes
        else:
            if random.random() < self.p:
                h_img, _, _ = img.shape
                img = img[::-1, :, :]
                bboxes[:, [1, 3]] = h_img-bboxes[:, [1, 3]]
                bboxes = bboxes[:,[0,3,2,1,4]]
            return img, bboxes

class RandomCrop(object):
    '''随机裁切'''
    def __init__(self, use=True):
        self.use = use
        self.p = 0.5 if self.use else 1

    def __call__(self, img, bboxes):
        if self.use==False:
            return img, bboxes
        else:
            if random.random() < self.p:
                h_img, w_img, _ = img.shape

                max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)
                max_l_trans = max_bbox[0]
                max_u_trans = max_bbox[1]
                max_r_trans = w_img - max_bbox[2]
                max_d_trans = h_img - max_bbox[3]

                crop_xmin = max(0, int(max_bbox[0] - random.uniform(0, max_l_trans)))
                crop_ymin = max(0, int(max_bbox[1] - random.uniform(0, max_u_trans)))
                crop_xmax = max(w_img, int(max_bbox[2] + random.uniform(0, max_r_trans)))
                crop_ymax = max(h_img, int(max_bbox[3] + random.uniform(0, max_d_trans)))

                img = img[crop_ymin : crop_ymax, crop_xmin : crop_xmax]

                bboxes[:, [0, 2]] = bboxes[:, [0, 2]] - crop_xmin
                bboxes[:, [1, 3]] = bboxes[:, [1, 3]] - crop_ymin
            return img, bboxes

class Blur(object):
    '''添加模糊，输入值blur=1是让图片明显模糊，blur=2让图片出现小的模糊'''
    def __init__(self, use=True, blur=2):
        self.use = use
        self.p = 0.5 if self.use else 1
        self.blur = blur

    def __call__(self, img, bboxes):
        if self.use==False:
            return img, bboxes
        else:
            if random.random() < self.p:
                if self.blur == 1:
                    img = cv2.GaussianBlur(img, (17, 17), 0)        # 大模糊
                else:
                    ksize = (self.blur / 2) * 2 + 1
                    img = cv2.GaussianBlur(img, (int(ksize), int(ksize)), 0)    # 小模糊

            return img,bboxes

class Gaussian_Noise(object):
    def __init__(self, use=True, blur=2):
        self.use = use
        self.p = 0.2 if self.use else 1
        self.blur = blur

    def __call__(self, img, bboxes):
        if self.use==False:
            return img, bboxes
        else:
            if random.random() < self.p:
                noise = np.array(img.shape)
                gaussian_noise = np.random.randint(0, 127)
                cv2.randn(noise, 0, gaussian_noise)  # mean and variance
                img = np.array(img + noise,dtype=np.uint8)


            return img,bboxes

class Random_Resize(object):
    '''保持竖屏或者横屏，但是宽高比会略微的变化'''
    def __init__(self, use =True,change= 0.9):
        self.p =0.5
        self.use = use
        self.change = change

    def get_new_hw(self,h,w):
        scale_w = np.random.randint(7,13)/10
        scale_h = np.random.randint(7,13)/10
        nw = int(w * scale_w)
        nh = int(h * scale_h)
        new_hw = nh/nw

        return nh,nw,new_hw

    def __call__(self, img, bboxes):
        if self.use==False:
            return img,bboxes
        else:
            if random.random()<self.p:
                h, w, _ = img.shape
                old_hw = h / w
                while(1):
                    nh, nw, new_hw = self.get_new_hw(h,w)
                    if new_hw < old_hw + 0.3 and new_hw > old_hw - 0.3:
                        break
                image = cv2.resize(img,(nw, nh))
                resize_ratio_w = nw / w
                resize_ratio_h = nh / h
                bbox =[]
                for box in bboxes:
                    box[0] = box[0] * resize_ratio_w
                    box[1] = box[1] * resize_ratio_h
                    box[2] = box[2] * resize_ratio_w
                    box[3] = box[3] * resize_ratio_h
                    if box[0]>nw or box[1]>nh:
                        continue
                    if nw-box[0]<=5 or nh-box[1]<=5:
                        continue
                    if box[2]>nw:
                        box[2]=nw
                    if box[3]>nh:
                        box[3]=nh
                    bbox.append(box)
                return image,bbox
            else:
                return img, bboxes

class Resize_to_shape(object):
    '''保持宽高比，保持一边为416'''
    def __init__(self,target_shape):
        self.input_shape = min(target_shape[0],target_shape[1])

    def __call__(self, img, bboxes):
        target_w = self.input_shape
        ih, iw, _ = img.shape
        new_ar = iw / ih  # >2.5 or <0.4
        if new_ar >= 1:
            nw = target_w
            nh = nw * (ih / iw)
        else:
            nh = target_w
            nw = nh * (iw / ih)
        scale_w = nw / iw
        scale_h = nh / ih
        if bboxes!= []:
            bboxes[:, [0, 2]] = scale_w * bboxes[:, [0, 2]]
            bboxes[:, [1, 3]] = scale_h * bboxes[:, [1, 3]]
        image_resized = cv2.resize(img, (int(nw), int(nh)))
        return image_resized, bboxes

class Resize(object):
    """
    Resize the image to target size and transforms it into a color channel(BGR->RGB),
    as well as pixel value normalization([0,1])
    """
    def __init__(self, target_shape, correct_box=False):
        self.h_target, self.w_target = target_shape
        self.correct_box = correct_box

    def __call__(self, img, bboxes):
        h_org , w_org , _= img.shape

        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)
        resize_ratio = min(1.0 * self.w_target / w_org, 1.0 * self.h_target / h_org)
        if self.correct_box:
            scale = np.random.randint(7,9)/10
        else:
            scale = 1
        resize_w = int(resize_ratio * w_org * scale)
        resize_h = int(resize_ratio * h_org * scale)
        image_resized = cv2.resize(img, (resize_w, resize_h),interpolation=cv2.INTER_AREA)

        image_paded = np.full((self.h_target, self.w_target, 3), 128.0)
        image_paded[:resize_h, :resize_w, :] = image_resized

        # image = image_paded / 255.0  # normalize to [0, 1]
        image = (image_paded).astype(np.uint8)

        if len(bboxes)==0:
            return image, []
        else:
            bboxes[:, [0, 2]] = bboxes[:, [0, 2]] * resize_ratio * scale
            bboxes[:, [1, 3]] = bboxes[:, [1, 3]] * resize_ratio * scale
            return image, bboxes

class Affine_Transformation(object):
    '''仿射变换'''
    def __init__(self,):
        pass

    def __call__(self, img, bboxes):
        ih, iw, _ = img.shape

        # 平移变换
        C = np.eye(3)
        C[0,2] = 0
        C[1,2] = 0

        # 缩放
        V = np.eye(3)
        if random.randint(0, 1) == 1:
            V[0,0] = 1
            V[1,1] = 1

        # 左右翻转
        N = np.eye(3)
        if random.randint(0,1)==1:
            print('chance')
            N[0,0] = -1
            N[0,2] = iw
            flag_left_right = True
        else:
            flag_left_right = False

        # 上下翻转
        M = np.eye(3)
        if random.randint(0, 1) == 1:
            M[1,1] = -1
            M[1,2] = ih
            flag_up_down = True
        else:
            flag_up_down = False

        #剪切变换
        S = np.eye(3)
        if random.randint(0, 1) == 1:
            shear = 0       # 10
            S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)
            S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)

        # 旋转和缩放
        R = np.eye(3)
        if random.randint(0, 1) == 1:
            degrees=0
            scale = 0       # 0.1
            a = random.uniform(-degrees, degrees)
            # a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations
            s = random.uniform(1 - scale, 1 + scale)
            # s = 2 ** random.uniform(-scale, scale)
            R[:2] = cv2.getRotationMatrix2D(angle=a, center=(416/2, 416/2), scale=s)  # angle=旋转角度，center=旋转中心 scale=缩放比例

        M = C@V@S@R@N@M
        img=cv2.warpAffine(img,M[:2],dsize=(416,416),borderValue=(114, 114, 114))

        xy = np.ones((len(bboxes)*2,3))
        a = bboxes[:,:4].reshape(-1,2)
        xy[:,:2] = a
        temp = xy@ M.T
        bboxes[:,:4] = temp[:,:2].reshape(-1,4)
        if flag_up_down==True:
            bboxes = bboxes[:,[0,3,2,1,4]]
        if flag_left_right == True:
            bboxes = bboxes[:, [2, 1, 0, 3, 4]]
        return img, bboxes

def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), bboxes=None,auto=False, scaleFill=False, scaleup=True):
    # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
    scaleup = False     # 只会进行缩小，不会放大
    auto = False         # 自动将图片尽量保持比例的，一边缩小为416，而且会添加一点点的padding
    scaleFill = False    # 自动缩放，让图片占满整个（416,416）,不留任何的padding
    shape = img.shape[:2]  # current shape [height, width]
    print(shape)
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better test mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding
    elif scaleFill:  # stretch  伸展
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    print(img.shape)
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border

    print('*'*50)

    bboxes[:, 0] = ratio[0]*bboxes[:,0] + dw  # pad width
    bboxes[:, 1] = ratio[1]*bboxes[:,1] + dh  # pad height
    bboxes[:, 2] = ratio[0]*bboxes[:,2] + dw
    bboxes[:, 3] = ratio[1]*bboxes[:,3] + dh
    return img, bboxes


if __name__ == '__main__':
    import config.Yolov4 as cfg
    img_txt_path = cfg.annotation_path
    with open(img_txt_path) as f:
        annotations = list(filter(lambda x: len(x) > 0, f.readlines()))

    def get_img_and_box(index):
        i = np.random.randint(0, len(annotations))
        anno = annotations[i].strip().split(' ')
        img_path = anno[0]
        bboxes = np.array([list(map(float, box.split(','))) for box in anno[1:]])
        img = cv2.imread(img_path)
        print(img_path)
        return img,bboxes

    class2id = ['heel', 'maid']
    for b in range(10):
        img,bboxes = get_img_and_box(b)
        # img, bboxes = RandomHorizontalFilp()(np.copy(img), np.copy(bboxes))
        # img,bboxes = Resize([416,416])(np.copy(img), np.copy(bboxes))
        img,bboxes = Affine_Transformation()(np.copy(img), np.copy(bboxes))

        # img, bboxes = letterbox(img,(416,416),bboxes=bboxes)

        if random.random() < 0:
            img2, bboxes2 = get_img_and_box(b)
            img2, bboxes2 = Resize([416,416])(np.copy(img2), np.copy(bboxes2))
            r = np.random.beta(8.0, 8.0)  # mixup ratio, alpha=beta=8.0
            img = (img * r + img2 * (1 - r)).astype(np.uint8)
            bboxes = np.concatenate((bboxes, bboxes2), 0)

        img = np.array(img, dtype=np.uint8)
        h, w, c = img.shape
        tl = round(0.002 * max(img.shape[0:2])) + 1  # line thickness
        color = [random.randint(0, 255) for _ in range(3)]
        for i in bboxes:
            c1,c2 = (round(i[0]),round(i[1])), (round(i[2]),round(i[3]))
            cv2.rectangle(img,c1,c2,color,thickness=tl)
            label = class2id[int(i[4])]
            if label:
                tf = max(tl - 1, 1)  # font thickness
                t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
                c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
                cv2.rectangle(img, c1, c2, color, -1)  # filled
                cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [0, 0, 0], thickness=tf, lineType=cv2.LINE_AA)
        cv2.imshow('change',img)
        cv2.waitKey(0)